Write a custom CUDA kernel to optimize `Sparsemax Loss` (ICML 2016).

Formula: L = 0.5 * sum_{j in Support} (z_j^2 - tau^2) + 0.5 - z_target
Algorithm to find Support and tau:
1. Sort logits z in descending order.
2. Find largest k such that 1 + k * z_k > sum(z_1...z_k).
3. tau = (sum(z_1...z_k) - 1) / k.
4. Support set is indices where z_j > tau.

Problem Analysis:
1. Sorting Overhead: The standard implementation uses `torch.sort`, which operates in global memory and is expensive for the subsequent logic flow.
2. Memory Traffic: Calculating cumsum and masks after sorting requires multiple passes over global memory tensors.

Optimization Strategy: Fused Shared-Memory Sort & Reduction

Constraint: Assume `num_classes` is a power of 2 (e.g., 2048) to facilitate efficient Bitonic Sort.

1. Block-per-Row: Launch one block per sample.
2. Shared Memory Loading: Load the entire row of logits into Shared Memory.
3. Bitonic Sort (Descending): Implement parallel Bitonic Sort in Shared Memory to order the logits. This avoids global memory sorting.
4. Parallel Scan (Cumsum): Compute the prefix sum of the sorted logits in Shared Memory to evaluate the condition `1 + k * z_k > cumsum_k`.
5. Threshold Detection: Identify the threshold index `k` and compute `tau`.
6. Fused Loss Calculation:
   - Calculate sum of squares for the top-k elements (using reduction).
   - Calculate final loss using the pre-loaded target logit (read from global memory initially). 
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn

BATCH_SIZE = 2048
NUM_CLASSES = 2048 
SHAPE = (BATCH_SIZE, NUM_CLASSES)

class SparsemaxLoss(nn.Module):
    """
    Sparsemax Loss (Martins & Astudillo, 2016)
    L = 0.5 * sum(z_j^2 - tau^2) + 0.5 - z_y
    """
    def __init__(self, reduction='mean'):
        super(SparsemaxLoss, self).__init__()
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # targets: (N)
        
        # Sort (Descending) 
        z_sorted, _ = torch.sort(logits, dim=1, descending=True)
        
        z_cumsum = torch.cumsum(z_sorted, dim=1)
        
        k = torch.arange(1, logits.size(1) + 1, device=logits.device)
        
        support = (1 + k * z_sorted) > z_cumsum
        k_z = torch.sum(support, dim=1, keepdim=True) # (N, 1)
        
        zs_sum = torch.gather(z_cumsum, 1, k_z - 1)
        tau = (zs_sum - 1) / k_z
        
        mask = torch.arange(NUM_CLASSES, device=logits.device).unsqueeze(0) < k_z
        z_support = z_sorted * mask
        
        sum_sq_z = (z_support ** 2).sum(dim=1)
        sum_sq_tau = (tau.squeeze(1) ** 2) * k_z.squeeze(1).float()
        
        z_y = logits.gather(1, targets.unsqueeze(1)).squeeze(1)
        
        loss = 0.5 * (sum_sq_z - sum_sq_tau) + 0.5 - z_y
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = SparsemaxLoss(reduction=reduction)
    
    def forward(self, logits, targets):
        return self.loss_fn(logits, targets)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), targets.contiguous()]

def get_init_inputs():
    return ['none']